import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import logging
from data_fetcher import DataFetcher
from data_preprocessor import DataPreprocessor
from strategies import PriceVolumeStrategy, FinancialStrategy, MLStrategy
from backtest import Backtest

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def main():
    # 设置Tushare Pro token
    token = "your_tushare_token"  # 请替换为你的Tushare Pro token
    
    # 数据获取和预处理
    data_fetcher = DataFetcher(token)
    data_preprocessor = DataPreprocessor()
    
    # 回测参数
    start_date = "20200101"
    end_date = "20231231"
    initial_capital = 1000000  # 初始资金100万
    
    # 获取股票列表
    stocks = data_fetcher.get_stock_list()
    
    # 选择测试的股票
    test_stocks = stocks.head(5)['ts_code'].tolist()
    logger.info(f"选择的测试股票: {test_stocks}")
    
    # 创建回测系统
    backtest = Backtest(initial_capital=initial_capital)
    
    # 初始化策略
    price_volume_strategy = PriceVolumeStrategy()
    financial_strategy = FinancialStrategy()
    ml_strategy = MLStrategy()
    
    # 用于存储所有股票的回测结果
    all_results = []
    
    # 对每支股票进行回测
    for ts_code in test_stocks:
        logger.info(f"开始回测股票: {ts_code}")
        
        # 获取日频数据
        daily_data = data_fetcher.get_daily_data(ts_code, start_date, end_date)
        if daily_data.empty:
            logger.warning(f"{ts_code}日频数据为空，跳过")
            continue
        
        # 预处理日频数据
        daily_data = data_preprocessor.preprocess_daily_data(daily_data)
        
        # 获取财务数据
        financial_data = data_fetcher.get_financial_indicators(ts_code, start_date, end_date)
        
        # 合并财务数据
        merged_data = data_preprocessor.preprocess_financial_data(financial_data, daily_data)
        
        # 获取基准数据
        benchmark_data = data_fetcher.get_hs300_data(start_date, end_date)
        if not benchmark_data.empty:
            benchmark_data = data_preprocessor.preprocess_daily_data(benchmark_data)
        
        # 训练机器学习模型
        ml_strategy.train(merged_data.copy())
        
        # 为每种策略生成信号
        pv_signals = price_volume_strategy.generate_signals(merged_data.copy())
        fs_signals = financial_strategy.generate_signals(merged_data.copy())
        ml_signals = ml_strategy.generate_signals(merged_data.copy())
        
        # 运行回测
        pv_results = backtest.run(pv_signals, price_volume_strategy)
        fs_results = backtest.run(fs_signals, financial_strategy)
        ml_results = backtest.run(ml_signals, ml_strategy)
        
        # 评估回测结果
        pv_metrics = backtest.evaluate(pv_results, price_volume_strategy, benchmark_data)
        fs_metrics = backtest.evaluate(fs_results, financial_strategy, benchmark_data)
        ml_metrics = backtest.evaluate(ml_results, ml_strategy, benchmark_data)
        
        # 添加股票代码到评估结果
        pv_metrics['ts_code'] = ts_code
        fs_metrics['ts_code'] = ts_code
        ml_metrics['ts_code'] = ts_code
        
        # 保存结果
        all_results.extend([pv_metrics, fs_metrics, ml_metrics])
        
        # 绘制回测结果图表
        backtest.plot_results(pv_results, price_volume_strategy, benchmark_data, f"pv_{ts_code}.png")
        backtest.plot_results(fs_results, financial_strategy, benchmark_data, f"fs_{ts_code}.png")
        backtest.plot_results(ml_results, ml_strategy, benchmark_data, f"ml_{ts_code}.png")
    
    # 汇总所有回测结果
    if all_results:
        results_df = pd.DataFrame(all_results)
        
        # 保存结果到CSV文件
        results_df.to_csv("strategy_comparison.csv", index=False)
        logger.info("所有回测结果已保存至strategy_comparison.csv")
        
        # 计算每种策略的平均表现
        strategy_avg = results_df.groupby('strategy').mean()
        logger.info("各策略平均表现:\n" + str(strategy_avg))
        
        # 可视化各策略对比
        plot_strategy_comparison(strategy_avg)

def plot_strategy_comparison(strategy_avg):
    """
    可视化各策略对比
    :param strategy_avg: 各策略平均表现DataFrame
    """
    try:
        metrics = ['annual_return', 'max_drawdown', 'sharpe_ratio', 'win_rate', 'profit_loss_ratio']
        metric_names = ['年化收益率', '最大回撤', '夏普比率', '胜率', '盈亏比']
        
        plt.figure(figsize=(15, 10))
        
        for i, metric in enumerate(metrics):
            plt.subplot(2, 3, i+1)
            bars = plt.bar(strategy_avg.index, strategy_avg[metric])
            
            if metric == 'max_drawdown':
                plt.bar_label(bars, padding=3, fmt='%.2f%%')
            elif metric in ['annual_return', 'win_rate']:
                plt.bar_label(bars, padding=3, fmt='%.2f%%', label_type='center')
            else:
                plt.bar_label(bars, padding=3, fmt='%.2f', label_type='center')
            
            plt.title(metric_names[i])
            plt.xticks(rotation=45)
        
        plt.tight_layout()
        plt.savefig("strategy_comparison.png")
        logger.info("策略对比图表已保存至strategy_comparison.png")
        plt.show()
    except Exception as e:
        logger.error(f"策略对比可视化失败: {e}")

if __name__ == "__main__":
    main()    